#include "Containers/RangePair.h"
#include "File/BOSSOutputLoader.h"
#include "File/YamlLoader.h"
#include "Fit/RooFitBuilder.h"
#include "Functions/Draw.h"
#include "Functions/Error.h"
#include "Functions/File.h"
#include "Functions/PrintTable.h"
#include "Functions/TerminalIO.h"
#include "Globals/Settings.h"
#include "Particle/ReconstructedParticle.h"
#include "RooFitResult.h"
#include "RooHistPdf.h"
#include "RooPlot.h"
#include "TArrow.h"
#include "TCanvas.h"
#include "TLatex.h"
#include "TLine.h"
#include <cmath>
#include <iostream>

using namespace BOSS_Afterburner;
using namespace BOSS_Afterburner::Draw;
using namespace BOSS_Afterburner::Error;
using namespace BOSS_Afterburner::File;
using namespace BOSS_Afterburner::TerminalIO;
using namespace RooFit;
using namespace BOSS_Afterburner::Detector::Units;
using namespace std;

Double_t scale = 1.;

class LikelihoodScan
{
public:
  LikelihoodScan(const YAML::Node& conf);
  Double_t CL       = .9;
  Double_t step     = 1.;
  Double_t max      = 250.;
  string   logscale = "";
  Bool_t   loaded   = false;
  Bool_t   debug    = false;
  Bool_t   print    = false;
  Bool_t   write    = false;
};

LikelihoodScan::LikelihoodScan(const YAML::Node& conf)
{
  const auto& llScanDef = conf["Likelihood scan"];
  if(!llScanDef || !llScanDef.IsMap()) return;
  const auto& clDef    = llScanDef["confidence level"];
  const auto& stepDef  = llScanDef["step size"];
  const auto& maxDef   = llScanDef["max"];
  const auto& scaleDef = llScanDef["scale"];
  const auto& logDef   = llScanDef["log scale"];
  const auto& debugDef = llScanDef["print debug"];
  const auto& printDef = llScanDef["print table"];
  const auto& writeDef = llScanDef["write table"];
  if(clDef.IsDefined() && clDef.IsScalar()) CL = clDef.as<Double_t>();
  if(stepDef.IsDefined() && stepDef.IsScalar()) step = stepDef.as<Double_t>();
  if(maxDef.IsDefined() && maxDef.IsScalar()) max = maxDef.as<Double_t>();
  if(scaleDef.IsDefined() && scaleDef.IsScalar()) scale = scaleDef.as<Double_t>();
  if(logDef.IsDefined() && logDef.IsScalar()) logscale = logDef.as<string>();
  if(debugDef.IsDefined() && debugDef.IsScalar()) debug = debugDef.as<Bool_t>();
  if(printDef.IsDefined() && printDef.IsScalar()) print = printDef.as<Bool_t>();
  if(writeDef.IsDefined() && writeDef.IsScalar()) write = writeDef.as<Bool_t>();
  loaded = true;
}

void PlotDataHistOnFrame(RooDataHist& hist, RooPlot* frame)
{
  hist.plotOn(frame, LineWidth(2), LineColor(kBlue + 2), LineWidth(1), MarkerColor(kBlue + 2),
              MarkerSize(.5));
}

void SaveFrame(RooPlot* frame, const char* name)
{
  TCanvas c;
  frame->Draw();
  SaveCanvas(name, &c);
}

double LikelihoodExp(double logL)
{
  return exp(-scale * logL);
}

void ExclToData(const char* config_file)
{
  YamlLoader       conf(config_file);
  BOSSOutputLoader file(conf);

  // * Get particle and draw options * //
  ReconstructedParticle particle(conf["Reconstructed particle"]);
  BranchPlotOptions     drawOpts(conf["Plot to analyse"]);

  auto        signalHist = file.Draw(drawOpts);
  const char* unit       = drawOpts[0].GetUnit();

  RangePair fit  = conf["Ranges"]["fit"];
  RangePair plot = conf["Ranges"]["plot"];

  // * Construct data set + variable * //
  const char* varname =
    Form("#it{M}_{%s #rightarrow %s}", particle.NameLaTeX(), particle.DaughterLabel());
  RooRealVar    var(varname, varname, drawOpts[0].From(), drawOpts[0].To(), unit);
  RooFitBuilder builder(var);

  // * Construct RooFit data histogram * //
  const auto& dataDef = conf["Plot to analyse"];
  if(!dataDef || !dataDef.IsMap()) EXCEPTION("Missing key \"Plot to analyse\"");
  auto        hist = file.Draw(dataDef);
  RooDataHist dataHist("data", hist->GetTitle(), var, Import(*hist));

  // * Build complete fit model * //
  const auto& modelDef = conf["Fit model"];
  if(!modelDef || !modelDef.IsMap()) EXCEPTION("Missing key \"Fit model\"");
  SigPlusBck model = builder.BuildSigPlusBck(modelDef);

  // * Plot data * //
  RooPlot* frame = var.frame();
  dataHist.plotOn(frame, LineWidth(2), LineColor(kBlue + 2), LineWidth(1), MarkerColor(kBlue + 2),
                  MarkerSize(.5));
  PlotStatsOnFrame(dataHist, frame);

  // * Legend position * //
  Double_t xMin     = .68;
  Double_t xMax     = .97;
  Double_t yMin     = .92;
  Double_t textsize = .033;

  // * Perform fit convoluted with truncated polynomial background * //
  Color_t colour;
  colour      = kRed;
  auto result = model.pdf->fitTo(dataHist, Range(fit.X1(), fit.X2()));
  model.pdf->plotOn(frame, LineWidth(2), LineColor(colour), Range(plot.X1(), plot.X2()),
                    LineStyle(kDotted));
  model.pdf->plotOn(frame, LineWidth(2), LineColor(colour));
  if(model.bck.pdf)
    model.pdf->plotOn(frame, Components(*model.bck.pdf), LineStyle(kDotted), LineWidth(2),
                      LineColor(kBlack));

  model.pdf->paramOn(frame, Layout(xMin, xMax, yMin), Format("NU", AutoPrecision(3)));
  if(frame)
  {
    frame->getAttText()->SetTextSize(textsize);
    frame->getAttText()->SetTextColor(colour);
  }

  TCanvas c;
  c.SetBatch();
  frame->Draw();

  // * Get parameters * //
  RooArgSet data;
  auto      params = model.pdf->getParameters(data);

  // * Draw indications * //
  const auto& srDef = conf["3-sigma range"];
  try
  {
    Indicate3SigmaRange(conf);
  }
  catch(const YAML::TypedBadConversion<Double_t>& e)
  {
    double width = 0.;
    for(const auto& sig : srDef["sigmas"])
    {
      auto par = dynamic_cast<RooRealVar*>(params->find(sig.as<string>().c_str()));
      if(par) width += par->getVal();
    }
    auto par = dynamic_cast<RooRealVar*>(params->find(srDef["mass"].as<string>().c_str()));
    if(!par) throw YAML::BadConversion(YAML::Mark::null_mark());
    Indicate3SigmaRange(par->getVal(), width);
  }
  IndicateFitRange(conf);
  IndicateMassWindow(conf);

  // * Draw and write fitted histograms * //
  TString outputName;
  if(model.sig.pdf) outputName += model.sig.name + "_";
  if(model.bck.pdf) outputName += model.bck.name + "_";
  outputName.ReplaceAll(" ", "");
  SaveCanvas(Form("%s%s", outputName.Data(), particle.Name()), &c, drawOpts.LogXYZ().c_str());

  // * Fit results * //
  TerminalIO::PrintHeader("Fit results for first fit");
  auto     NsigName = conf["Labels"]["Nsig"].as<string>();
  auto     NbckName = conf["Labels"]["Nbck"].as<string>();
  auto     Nsig     = dynamic_cast<RooRealVar*>(params->find(NsigName.c_str()));
  auto     Nbck     = dynamic_cast<RooRealVar*>(params->find(NbckName.c_str()));
  Double_t nsig     = Nsig->getVal();
  Double_t nbck     = Nbck->getVal();
  cerr << "  Signal:           N_{sig} = " << nsig << endl;
  cerr << "  Background:       N_{bck} = " << nbck << endl;
  cerr << "  Chi square:  \\chi^2_{red} = " << frame->chiSquare(2) << endl;
  // cerr << "  Likelihood:       -log(L) = " << result->minNll() << endl;
  // cerr << "                         L  = " << exp(-0.5 * result->minNll()) << endl;

  // * Compute likelihoods * //
  LikelihoodScan ll(conf);
  if(!ll.loaded) return;
  Int_t    nBins    = floor(ll.max / ll.step) + 1;
  Double_t binWidth = ll.step;

  TerminalIO::PrintHeader("Likelihood scan");
  TH1D llHist("likelihood", "likelihood", nBins, -binWidth * 0.5,
              nBins * binWidth + binWidth * 0.5);
  llHist.SetTitle("Likelihood scan");
  llHist.GetXaxis()->SetTitle("#it{N}_{sig} #times N_{y}");
  llHist.GetYaxis()->SetTitle("L (a.u.)");
  llHist.GetYaxis()->SetLabelSize(0.);
  llHist.GetYaxis()->SetNdivisions(1);
  llHist.SetLineColor(kRed);
  llHist.SetLineWidth(3);

  RooMsgService::instance().setSilentMode(true);
  Table    table;
  Double_t firstval;
  for(Int_t i = 0; i < nBins; ++i)
  {
    double     nsig_new = i * binWidth;
    RooRealVar Nsig_L("#it{N}_{S_{new}}", "#it{N}_{S_{new}}", nsig_new);
    RooArgList coeff_L(Nsig_L, *Nbck);
    auto       shape_L  = builder.AddPdfs(model.sum, &coeff_L);
    auto       result_L = shape_L->fitTo(dataHist, Range(fit.X1(), fit.X2()), Save(true));
    auto       frame_L  = var.frame();
    PlotDataHistOnFrame(dataHist, frame_L);
    shape_L->plotOn(frame_L);

    double L     = result_L->minNll();
    double chisq = frame_L->chiSquare(2);
    double val   = LikelihoodExp(L);
    if(!i) firstval = val;
    delete shape_L, result_L, frame_L;
    llHist.SetBinContent(i + 1, val / firstval);
    if(ll.debug) cerr << nsig_new << "\t" << L << "\t" << val << "\t" << chisq << endl;
    table << nsig_new << L << val << chisq << TableRow::endrow;
  }
  table.SetHeader({"N_{sig}", "-log(L)", "L", "\\chi^2"});
  if(ll.print) table.Print();
  if(ll.write)
  {
    table.Write("LikelihoodScan");
    table.WriteLaTeX("LikelihoodScan");
  }

  // * Compute upper limit * //
  Double_t* cumulate = llHist.GetIntegral();
  Double_t  UL       = 0.;
  Int_t     bin_UL      = 0;
  for(Int_t i = 0; i < nBins; i++)
  {
    UL  = llHist.GetBinCenter(i + 1);
    bin_UL = i;
    if(cumulate[i] > ll.CL) break;
  }
  PrintBold("Resulting upper limit:");
  cerr << "  N_{sig} = " << nsig << endl;
  cerr << "  N_{UL}  = " << UL << endl;

  // * Integral hist * //
  TH1D integral(llHist);
  for(Int_t i = bin_UL + 1; i < nBins; ++i)
    integral.SetBinContent(i + 1, 0.);
  integral.SetLineWidth(0.);
  integral.SetFillColor(kGray);
  integral.SetFillStyle(1001);

  // * Draw likelihood * //
  TCanvas c2;
  llHist.Draw("c");
  integral.Draw("same");
  llHist.Draw("c same");
  c2.Update();
  Double_t cx1, cy1, cx2, cy2;
  c2.GetRangeAxis(cx1, cy1, cx2, cy2);
  Double_t scaleX = cx2 - cx1;
  Double_t scaleY = cy2 - cy1;

  Double_t y = llHist.GetBinContent(bin_UL + 1);

  // * Nsig indication * //
  TLine  linesig(nsig, 0., nsig, llHist.GetBinContent(llHist.GetMaximumBin()));
  TLatex textsig(nsig, .6 * llHist.GetBinContent(llHist.GetMaximumBin()),
                 Form("#it{N}_{sig} = %g", nsig));
  textsig.SetTextSize(.04);
  textsig.SetTextAngle(-90.);
  textsig.SetTextAlign(21);
  linesig.SetLineStyle(kDotted);

  // * UL arrow * //
  TArrow arrow(UL, y + .18 * scaleY, UL, y + .03 * scaleY);
  arrow.SetArrowSize(.02);
  arrow.SetLineWidth(2.);

  // * UL text * //
  TLatex textUL(
    UL, y + .2 * scaleY,
    Form("#it{N}_{UL} = %g", UL, bin_UL * binWidth));
  textUL.SetTextSize(.045);

  // * C.L. text * //
  TLatex text90(.5 * UL, .3 * llHist.GetBinContent(llHist.FindBin(.5 * UL)),
                Form("%g%%", ll.CL * 100.));
  text90.SetTextAlign(22);
  text90.SetTextColor(kGray + 2);
  text90.SetTextSize(.07);

  // * Draw indications * //
  arrow.Draw();
  textUL.Draw();
  text90.Draw();
  textsig.Draw();
  linesig.Draw();
  SaveCanvas("Likelihood", &c2, ll.logscale.c_str());
}

/// Main function that is called when executing the executable compiled using e.g. `g++`.
int main(int argc, char* argv[])
{
  TerminalIO::CheckMainArguments(argc, argv, 1);
  ExclToData(argv[1]);
  return 0;
}